from openai import OpenAI
import os

def chatNow(model='gpt-3.5-turbo-instruct',mode='balance'):
    #根据mode来调整temperature及presence_penalty

    if mode == 'balance':
        temperature = 1
        presence_penalty = 0
    elif mode == 'precision':
        temperature = 0.8
        presence_penalty = 2
    elif mode == 'creativity':
        temperature = 1.2
        presence_penalty = -1

    def chat(prompt):
        try:
            client = OpenAI(api_key="sk-yrpeg76LZt53LhyTcue8ojFcvRkzs53fapexPklNqR5KPKae", base_url="https://api.fe8.cn/v1")
            response = client.completions.create(model=model, prompt=prompt, max_tokens=150, temperature=temperature, presence_penalty=presence_penalty)

            return response.choices[0].text.strip()
        except Exception as e:
            return "broken"

    text= ""
    turns = []

    while True:
        userInput = input("请输入要提的问题?")
        if len(userInput) == 0:
            print("您必须输入问题内容！")
            continue
        elif userInput == "quit":
            print("所有对话结束。")
            break


        prompt = text + f"Human: {userInput}\nAI:"
        response = chat(prompt)
        if response == 'broken':
            retryTimes = 3
            while response == 'broken' and retryTimes>=0:
                retryTimes = retryTimes - 1
                response = chat(prompt)
        print(response)

        turns.append({"user": userInput, "assistant": response})

        if len(turns) <= 10:
            text = ""
            for turn in turns:
                text += f"Human: {turn['user']}\nAI: {turn['assistant']}\n"
        else:
            for turn in turns[-10:]:
                text += f"Human: {turn['user']}\nAI: {turn['assistant']}\n"

if __name__ == '__main__':
    chatNow(mode="creativity")

